/**
 * NetFlow - OpenWRT 网络流量管理系统
 * 带宽管理器模块
 * 
 * 本文件实现了带宽管理功能，
 * 包括静态带宽分配、动态带宽调整和自适应控制算法。
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <math.h>
#include <syslog.h>
#include <errno.h>
#include <time.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netlink/netlink.h>
#include <netlink/route/link.h>
#include <netlink/route/tc.h>
#include <netlink/route/qdisc.h>
#include <netlink/route/class.h>

#include "netflow.h"
#include "bandwidth_manager.h"
#include "priority_scheduler.h"
#include "util.h"

/* 全局变量 */
static char interface[IFNAMSIZ];
static uint32_t total_bandwidth;
static int if_index;
static struct nl_sock *nl_sock;
static network_metrics_t current_metrics;
static uint64_t last_update_time = 0;
static uint64_t last_rx_bytes = 0;
static uint64_t last_tx_bytes = 0;
static bandwidth_config_t bandwidth_config;

/**
 * 初始化带宽管理器
 */
int init_bandwidth_manager(const char *ifname, uint32_t bandwidth) {
    syslog(LOG_INFO, "初始化带宽管理器...");
    
    // 保存接口名称和带宽
    strncpy(interface, ifname, IFNAMSIZ - 1);
    total_bandwidth = bandwidth;
    
    // 创建netlink套接字
    nl_sock = nl_socket_alloc();
    if (!nl_sock) {
        syslog(LOG_ERR, "无法分配netlink套接字");
        return -1;
    }
    
    // 连接到netlink路由协议
    if (nl_connect(nl_sock, NETLINK_ROUTE) < 0) {
        syslog(LOG_ERR, "无法连接到netlink路由协议");
        nl_socket_free(nl_sock);
        return -1;
    }
    
    // 获取接口索引
    struct nl_cache *link_cache;
    if (rtnl_link_alloc_cache(nl_sock, AF_UNSPEC, &link_cache) < 0) {
        syslog(LOG_ERR, "无法获取链路缓存");
        nl_socket_free(nl_sock);
        return -1;
    }
    
    struct rtnl_link *link = rtnl_link_get_by_name(link_cache, interface);
    if (!link) {
        syslog(LOG_ERR, "无法找到接口: %s", interface);
        nl_cache_free(link_cache);
        nl_socket_free(nl_sock);
        return -1;
    }
    
    if_index = rtnl_link_get_ifindex(link);
    rtnl_link_put(link);
    nl_cache_free(link_cache);
    
    // 初始化网络指标
    memset(&current_metrics, 0, sizeof(network_metrics_t));
    current_metrics.traffic.capacity_bps = total_bandwidth * 1000; // 转换为bps
    
    // 初始化默认带宽配置
    init_default_bandwidth_config();
    
    // 初始化监控
    if (init_network_monitoring() < 0) {
        syslog(LOG_ERR, "无法初始化网络监控");
        nl_socket_free(nl_sock);
        return -1;
    }
    
    syslog(LOG_INFO, "带宽管理器初始化成功");
    return 0;
}

/**
 * 清理带宽管理器
 */
void cleanup_bandwidth_manager(void) {
    syslog(LOG_INFO, "清理带宽管理器...");
    
    // 释放netlink套接字
    if (nl_sock) {
        nl_socket_free(nl_sock);
        nl_sock = NULL;
    }
    
    syslog(LOG_INFO, "带宽管理器已清理");
}

/**
 * 初始化默认带宽配置
 */
void init_default_bandwidth_config(void) {
    // 低负载配置（负载 < 50%）
    bandwidth_config.low_load.min_bw[0] = 30; // 实时交互类
    bandwidth_config.low_load.min_bw[1] = 25; // 流媒体类
    bandwidth_config.low_load.min_bw[2] = 25; // 普通数据类
    bandwidth_config.low_load.min_bw[3] = 20; // 背景传输类
    
    bandwidth_config.low_load.max_bw[0] = 100; // 实时交互类
    bandwidth_config.low_load.max_bw[1] = 90;  // 流媒体类
    bandwidth_config.low_load.max_bw[2] = 80;  // 普通数据类
    bandwidth_config.low_load.max_bw[3] = 70;  // 背景传输类
    
    // 中等负载配置（50% <= 负载 < 80%）
    bandwidth_config.medium_load.min_bw[0] = 40; // 实时交互类
    bandwidth_config.medium_load.min_bw[1] = 30; // 流媒体类
    bandwidth_config.medium_load.min_bw[2] = 20; // 普通数据类
    bandwidth_config.medium_load.min_bw[3] = 10; // 背景传输类
    
    bandwidth_config.medium_load.max_bw[0] = 100; // 实时交互类
    bandwidth_config.medium_load.max_bw[1] = 80;  // 流媒体类
    bandwidth_config.medium_load.max_bw[2] = 60;  // 普通数据类
    bandwidth_config.medium_load.max_bw[3] = 40;  // 背景传输类
    
    // 高负载配置（负载 >= 80%）
    bandwidth_config.high_load.min_bw[0] = 50; // 实时交互类
    bandwidth_config.high_load.min_bw[1] = 30; // 流媒体类
    bandwidth_config.high_load.min_bw[2] = 15; // 普通数据类
    bandwidth_config.high_load.min_bw[3] = 5;  // 背景传输类
    
    bandwidth_config.high_load.max_bw[0] = 100; // 实时交互类
    bandwidth_config.high_load.max_bw[1] = 70;  // 流媒体类
    bandwidth_config.high_load.max_bw[2] = 40;  // 普通数据类
    bandwidth_config.high_load.max_bw[3] = 20;  // 背景传输类
}

/**
 * 初始化网络监控
 */
int init_network_monitoring(void) {
    // 读取初始接口统计信息
    if (read_interface_stats(&last_rx_bytes, &last_tx_bytes) < 0) {
        return -1;
    }
    
    last_update_time = current_time_ms();
    
    return 0;
}

/**
 * 读取接口统计信息
 */
int read_interface_stats(uint64_t *rx_bytes, uint64_t *tx_bytes) {
    char path[PATH_MAX];
    FILE *fp;
    
    // 读取接收字节数
    snprintf(path, sizeof(path), "/sys/class/net/%s/statistics/rx_bytes", interface);
    fp = fopen(path, "r");
    if (!fp) {
        syslog(LOG_ERR, "无法打开 %s: %s", path, strerror(errno));
        return -1;
    }
    
    if (fscanf(fp, "%lu", rx_bytes) != 1) {
        syslog(LOG_ERR, "无法读取接收字节数");
        fclose(fp);
        return -1;
    }
    
    fclose(fp);
    
    // 读取发送字节数
    snprintf(path, sizeof(path), "/sys/class/net/%s/statistics/tx_bytes", interface);
    fp = fopen(path, "r");
    if (!fp) {
        syslog(LOG_ERR, "无法打开 %s: %s", path, strerror(errno));
        return -1;
    }
    
    if (fscanf(fp, "%lu", tx_bytes) != 1) {
        syslog(LOG_ERR, "无法读取发送字节数");
        fclose(fp);
        return -1;
    }
    
    fclose(fp);
    
    return 0;
}

/**
 * 更新网络统计信息
 */
void update_network_statistics(void) {
    uint64_t rx_bytes, tx_bytes;
    uint64_t now = current_time_ms();
    uint64_t elapsed = now - last_update_time;
    
    if (elapsed < 1000) {
        // 至少等待1秒钟再更新
        return;
    }
    
    // 读取当前接口统计信息
    if (read_interface_stats(&rx_bytes, &tx_bytes) < 0) {
        return;
    }
    
    // 计算带宽使用率
    uint64_t total_bytes = (rx_bytes - last_rx_bytes) + (tx_bytes - last_tx_bytes);
    uint64_t bytes_per_sec = total_bytes * 1000 / elapsed;
    uint64_t bits_per_sec = bytes_per_sec * 8;
    
    // 更新网络指标
    current_metrics.traffic.current_bps = bits_per_sec;
    current_metrics.load_factor = (float)bits_per_sec / current_metrics.traffic.capacity_bps;
    
    // 更新队列统计信息
    update_queue_statistics();
    
    // 更新拥塞指标
    update_congestion_index();
    
    // 保存当前值用于下次计算
    last_rx_bytes = rx_bytes;
    last_tx_bytes = tx_bytes;
    last_update_time = now;
    
    // 记录统计信息
    if (current_metrics.load_factor > 0.9) {
        syslog(LOG_WARNING, "网络负载过高: %.2f%%", current_metrics.load_factor * 100);
    }
    
    // 根据负载情况，动态调整带宽分配
    if (current_metrics.load_factor > 0.8) {
        apply_bandwidth_config(&bandwidth_config.high_load);
    } else if (current_metrics.load_factor > 0.5) {
        apply_bandwidth_config(&bandwidth_config.medium_load);
    } else {
        apply_bandwidth_config(&bandwidth_config.low_load);
    }
}

/**
 * 更新队列统计信息
 */
void update_queue_statistics(void) {
    struct nl_cache *qdisc_cache;
    
    // 获取队列统计信息
    if (rtnl_qdisc_alloc_cache(nl_sock, if_index, &qdisc_cache) < 0) {
        syslog(LOG_WARNING, "无法获取队列规则缓存");
        return;
    }
    
    struct nl_object *obj;
    struct rtnl_qdisc *qdisc;
    
    // 遍历队列规则
    nl_object_for_each(obj, nl_cache_get_first(qdisc_cache)) {
        qdisc = (struct rtnl_qdisc *)obj;
        uint32_t handle = rtnl_tc_get_handle(TC_CAST(qdisc));
        uint32_t parent = rtnl_tc_get_parent(TC_CAST(qdisc));
        
        // 检查是否是我们的HTB子队列
        if (parent >= 0x10010 && parent <= 0x10040) {
            int prio_idx = (parent - 0x10010) / 0x10;
            
            // 获取队列统计信息
            current_metrics.queue.packets[prio_idx] = rtnl_qdisc_get_stat(qdisc, RTNL_TC_PACKETS);
            current_metrics.queue.bytes[prio_idx] = rtnl_qdisc_get_stat(qdisc, RTNL_TC_BYTES);
            current_metrics.queue.drops[prio_idx] = rtnl_qdisc_get_stat(qdisc, RTNL_TC_DROPS);
            current_metrics.queue.overlimits[prio_idx] = rtnl_qdisc_get_stat(qdisc, RTNL_TC_OVERLIMITS);
            
            // 队列长度统计
            current_metrics.queue.qlen[prio_idx] = rtnl_qdisc_get_stat(qdisc, RTNL_TC_QLEN);
        }
    }
    
    nl_cache_free(qdisc_cache);
}

/**
 * 更新拥塞指标
 */
void update_congestion_index(void) {
    // 计算综合拥塞指标
    float drop_factor = 0.0;
    float qlen_factor = 0.0;
    
    // 计算丢包率
    uint64_t total_packets = 0, total_drops = 0;
    for (int i = 0; i < 4; i++) {
        total_packets += current_metrics.queue.packets[i];
        total_drops += current_metrics.queue.drops[i];
    }
    
    if (total_packets > 0) {
        float drop_rate = (float)total_drops / total_packets;
        
        // 将丢包率映射到0-1范围
        if (drop_rate > 0.1) {
            drop_factor = 1.0; // 丢包率超过10%，表示严重拥塞
        } else {
            drop_factor = drop_rate * 10;
        }
    }
    
    // 计算队列长度因子
    uint32_t max_qlen = 0;
    for (int i = 0; i < 4; i++) {
        if (current_metrics.queue.qlen[i] > max_qlen) {
            max_qlen = current_metrics.queue.qlen[i];
        }
    }
    
    // 队列长度超过100表示拥塞
    if (max_qlen > 100) {
        qlen_factor = 1.0;
    } else {
        qlen_factor = max_qlen / 100.0;
    }
    
    // 综合负载、丢包和队列长度因子
    current_metrics.congestion_index = 0.6 * current_metrics.load_factor + 
                                       0.3 * drop_factor + 
                                       0.1 * qlen_factor;
    
    if (current_metrics.congestion_index > 0.8) {
        syslog(LOG_WARNING, "网络拥塞严重: %.2f", current_metrics.congestion_index);
    }
}

/**
 * 应用带宽配置
 */
void apply_bandwidth_config(bandwidth_allocation_t *bw_config) {
    // 计算实际带宽
    uint32_t min_bw[4], max_bw[4];
    
    for (int i = 0; i < 4; i++) {
        min_bw[i] = total_bandwidth * bw_config->min_bw[i] / 100;
        max_bw[i] = total_bandwidth * bw_config->max_bw[i] / 100;
    }
    
    // 更新HTB类参数
    for (int i = 0; i < 4; i++) {
        uint32_t classid = 0x10010 + (i * 0x10);
        update_htb_class(classid, min_bw[i], max_bw[i], i + 1);
    }
}

/**
 * 更新HTB类参数
 */
int update_htb_class(uint32_t classid, uint32_t min_rate, uint32_t max_rate, uint8_t priority) {
    struct rtnl_class *cls;
    int err;
    
    cls = rtnl_class_alloc();
    if (!cls) {
        syslog(LOG_ERR, "无法分配类");
        return -1;
    }
    
    rtnl_tc_set_link(TC_CAST(cls), if_index);
    rtnl_tc_set_parent(TC_CAST(cls), 0x10001);
    rtnl_tc_set_handle(TC_CAST(cls), classid);
    rtnl_class_set_kind(cls, "htb");
    
    // 设置类参数
    rtnl_htb_set_rate(cls, min_rate);
    rtnl_htb_set_ceil(cls, max_rate);
    rtnl_htb_set_prio(cls, priority);
    
    err = rtnl_class_add(nl_sock, cls, NLM_F_REPLACE);
    rtnl_class_put(cls);
    
    if (err < 0) {
        syslog(LOG_ERR, "无法更新HTB类 %x: %s", classid, nl_geterror(err));
        return -1;
    }
    
    return 0;
}

/**
 * 执行动态调整
 */
void perform_dynamic_adjustment(void) {
    // 计算调整因子
    float adjustment_factor = calculate_adjustment_factor(&current_metrics);
    
    // 调整带宽分配
    if (adjustment_factor > 0.8) {
        // 网络状况良好，放宽带宽限制
        apply_bandwidth_config(&bandwidth_config.low_load);
    } else if (adjustment_factor > 0.5) {
        // 网络状况一般，使用默认带宽分配
        apply_bandwidth_config(&bandwidth_config.medium_load);
    } else {
        // 网络状况较差，严格限制低优先级流量
        apply_bandwidth_config(&bandwidth_config.high_load);
    }
    
    // 调整队列参数
    adjust_queue_parameters(adjustment_factor);
}

/**
 * 计算调整因子
 */
float calculate_adjustment_factor(network_metrics_t *metrics) {
    float load_weight = 0.7;
    float congestion_weight = 0.3;
    
    float factor = load_weight * (1.0 - metrics->load_factor) + 
                  congestion_weight * (1.0 - metrics->congestion_index);
    
    // 限制调整因子范围
    if (factor < 0.1) {
        factor = 0.1;
    } else if (factor > 1.0) {
        factor = 1.0;
    }
    
    return factor;
}

/**
 * 调整队列参数
 */
void adjust_queue_parameters(float adjustment_factor) {
    // 根据调整因子调整队列长度
    uint32_t queue_length = calculate_queue_length(adjustment_factor);
    
    // 更新SFQ队列参数
    for (int i = 0; i < 4; i++) {
        uint32_t parent = 0x10010 + (i * 0x10);
        uint32_t handle = 0x20000 + (i * 0x10000);
        update_sfq_parameters(parent, handle, queue_length);
    }
    
    // 调整令牌桶参数
    adjust_token_bucket_parameters(adjustment_factor);
}

/**
 * 计算队列长度
 */
uint32_t calculate_queue_length(float adjustment_factor) {
    // 基础队列长度
    uint32_t base_length = 100;
    
    // 根据调整因子增加队列长度
    if (adjustment_factor > 0.8) {
        // 网络状况良好，增加队列长度
        return base_length * 1.5;
    } else if (adjustment_factor > 0.5) {
        // 网络状况一般，使用默认队列长度
        return base_length;
    } else {
        // 网络状况较差，减少队列长度
        return base_length * 0.5;
    }
}

/**
 * 更新SFQ队列参数
 */
int update_sfq_parameters(uint32_t parent, uint32_t handle, uint32_t limit) {
    struct rtnl_qdisc *qdisc;
    int err;
    
    qdisc = rtnl_qdisc_alloc();
    if (!qdisc) {
        syslog(LOG_ERR, "无法分配队列规则");
        return -1;
    }
    
    rtnl_tc_set_link(TC_CAST(qdisc), if_index);
    rtnl_tc_set_parent(TC_CAST(qdisc), parent);
    rtnl_tc_set_handle(TC_CAST(qdisc), handle);
    rtnl_qdisc_set_kind(qdisc, "sfq");
    
    // 设置SFQ参数
    rtnl_sfq_set_perturb(qdisc, 10);  // 每10秒重新哈希
    rtnl_sfq_set_quantum(qdisc, 1514);  // 最大数据包大小
    rtnl_sfq_set_limit(qdisc, limit);   // 队列长度
    
    err = rtnl_qdisc_add(nl_sock, qdisc, NLM_F_REPLACE);
    rtnl_qdisc_put(qdisc);
    
    if (err < 0) {
        syslog(LOG_ERR, "无法更新SFQ队列参数: %s", nl_geterror(err));
        return -1;
    }
    
    return 0;
}

/**
 * 调整令牌桶参数
 */
void adjust_token_bucket_parameters(float adjustment_factor) {
    // 根据调整因子调整令牌桶参数
    // HTB默认令牌桶大小为1.5*rate
    // 这里可以根据需要调整，但要谨慎，避免过度修改导致不稳定
    
    // 例如，可以调整令牌桶缓冲因子
    float buffer_factor;
    
    if (adjustment_factor > 0.8) {
        // 网络状况良好，增加缓冲
        buffer_factor = 2.0;
    } else if (adjustment_factor > 0.5) {
        // 网络状况一般，使用默认缓冲
        buffer_factor = 1.5;
    } else {
        // 网络状况较差，减少缓冲
        buffer_factor = 1.0;
    }
    
    // 这里需要HTB的更高级API来设置buffer/cbuffer参数
    // 但libnl可能不直接支持，需要使用tc命令或自定义的netlink消息
    // 简化起见，这里仅记录日志
    syslog(LOG_INFO, "调整令牌桶缓冲因子为 %.1f", buffer_factor);
}

/**
 * 获取当前网络指标
 */
void get_network_metrics(network_metrics_t *metrics) {
    if (metrics) {
        memcpy(metrics, &current_metrics, sizeof(network_metrics_t));
    }
}

/**
 * 设置带宽配置
 */
int set_bandwidth_config(bandwidth_config_t *config) {
    if (!config) {
        return -1;
    }
    
    // 检查配置有效性
    for (int i = 0; i < 4; i++) {
        // 检查百分比是否有效
        if (config->low_load.min_bw[i] > 100 || config->low_load.max_bw[i] > 100 ||
            config->medium_load.min_bw[i] > 100 || config->medium_load.max_bw[i] > 100 ||
            config->high_load.min_bw[i] > 100 || config->high_load.max_bw[i] > 100) {
            return -1;
        }
        
        // 检查最小值是否小于最大值
        if (config->low_load.min_bw[i] > config->low_load.max_bw[i] ||
            config->medium_load.min_bw[i] > config->medium_load.max_bw[i] ||
            config->high_load.min_bw[i] > config->high_load.max_bw[i]) {
            return -1;
        }
    }
    
    // 检查总和是否超过100%
    int low_min_sum = 0, med_min_sum = 0, high_min_sum = 0;
    for (int i = 0; i < 4; i++) {
        low_min_sum += config->low_load.min_bw[i];
        med_min_sum += config->medium_load.min_bw[i];
        high_min_sum += config->high_load.min_bw[i];
    }
    
    if (low_min_sum > 100 || med_min_sum > 100 || high_min_sum > 100) {
        return -1;
    }
    
    // 复制配置
    memcpy(&bandwidth_config, config, sizeof(bandwidth_config_t));
    
    // 应用新配置
    if (current_metrics.load_factor > 0.8) {
        apply_bandwidth_config(&bandwidth_config.high_load);
    } else if (current_metrics.load_factor > 0.5) {
        apply_bandwidth_config(&bandwidth_config.medium_load);
    } else {
        apply_bandwidth_config(&bandwidth_config.low_load);
    }
    
    return 0;
}

/**
 * 获取带宽配置
 */
void get_bandwidth_config(bandwidth_config_t *config) {
    if (config) {
        memcpy(config, &bandwidth_config, sizeof(bandwidth_config_t));
    }
}

/**
 * 设置总带宽
 */
int set_total_bandwidth(uint32_t bandwidth) {
    if (bandwidth < 1000 || bandwidth > 1000000) {
        // 带宽范围: 1Mbps - 1Gbps
        return -1;
    }
    
    total_bandwidth = bandwidth;
    current_metrics.traffic.capacity_bps = bandwidth * 1000; // 转换为bps
    
    // 应用新配置
    if (current_metrics.load_factor > 0.8) {
        apply_bandwidth_config(&bandwidth_config.high_load);
    } else if (current_metrics.load_factor > 0.5) {
        apply_bandwidth_config(&bandwidth_config.medium_load);
    } else {
        apply_bandwidth_config(&bandwidth_config.low_load);
    }
    
    return 0;
}

/**
 * 获取总带宽
 */
uint32_t get_total_bandwidth(void) {
    return total_bandwidth;
}

/**
 * 令牌桶流量整形
 */
int token_bucket_shaping(packet_info_t *packet, token_bucket_t *bucket) {
    uint64_t now = current_time_ms();
    
    // 计算上次操作后新产生的令牌数
    uint32_t new_tokens = (now - bucket->last_update) * bucket->rate / 1000;
    
    // 更新令牌桶中的令牌数，不超过桶容量
    bucket->tokens = (bucket->tokens + new_tokens) < bucket->capacity ? 
                     (bucket->tokens + new_tokens) : bucket->capacity;
    bucket->last_update = now;
    
    // 判断令牌是否足够
    if (bucket->tokens >= packet->size) {
        // 消耗令牌并允许数据包通过
        bucket->tokens -= packet->size;
        return 1;
    } else {
        // 令牌不足，延迟或丢弃数据包
        return 0;
    }
}

/**
 * 初始化令牌桶
 */
void init_token_bucket(